import numpy as np
from scipy.spatial.distance import directed_hausdorff
import torch

class SegmentationMetrics:
    """分割评估指标类"""
    
    def __init__(self, epsilon=1e-6):
        self.epsilon = epsilon
        
    def dice_coefficient(self, pred, target):
        """
        计算Dice系数
        
        Args:
            pred: 预测结果 (二值化)
            target: 真实标注
            
        Returns:
            dice: Dice系数 [0, 1]
        """
        pred = pred.astype(bool)
        target = target.astype(bool)
        
        intersection = np.logical_and(pred, target).sum()
        dice = 2.0 * intersection / (pred.sum() + target.sum() + self.epsilon)
        
        return dice
    
    def jaccard_index(self, pred, target):
        """
        计算Jaccard指数 (IoU)
        
        Args:
            pred: 预测结果
            target: 真实标注
            
        Returns:
            jaccard: Jaccard指数 [0, 1]
        """
        pred = pred.astype(bool)
        target = target.astype(bool)
        
        intersection = np.logical_and(pred, target).sum()
        union = np.logical_or(pred, target).sum()
        
        jaccard = intersection / (union + self.epsilon)
        
        return jaccard
    
    def positive_predictive_value(self, pred, target):
        """
        计算阳性预测值 (PPV/Precision)
        
        Args:
            pred: 预测结果
            target: 真实标注
            
        Returns:
            ppv: 阳性预测值 [0, 1]
        """
        pred = pred.astype(bool)
        target = target.astype(bool)
        
        true_positive = np.logical_and(pred, target).sum()
        predicted_positive = pred.sum()
        
        ppv = true_positive / (predicted_positive + self.epsilon)
        
        return ppv
    
    def sensitivity(self, pred, target):
        """
        计算敏感度 (Recall)
        
        Args:
            pred: 预测结果
            target: 真实标注
            
        Returns:
            sensitivity: 敏感度 [0, 1]
        """
        pred = pred.astype(bool)
        target = target.astype(bool)
        
        true_positive = np.logical_and(pred, target).sum()
        actual_positive = target.sum()
        
        sensitivity = true_positive / (actual_positive + self.epsilon)
        
        return sensitivity
    
    def hausdorff_distance(self, pred, target):
        """
        计算Hausdorff距离
        
        Args:
            pred: 预测结果
            target: 真实标注
            
        Returns:
            hd: Hausdorff距离
        """
        pred = pred.astype(bool)
        target = target.astype(bool)
        
        # 获取边界点
        pred_points = np.argwhere(pred)
        target_points = np.argwhere(target)
        
        if len(pred_points) == 0 or len(target_points) == 0:
            return np.inf
        
        # 计算双向Hausdorff距离
        forward_hd = directed_hausdorff(pred_points, target_points)[0]
        backward_hd = directed_hausdorff(target_points, pred_points)[0]
        
        hd = max(forward_hd, backward_hd)
        
        return hd
    
    def hausdorff_distance_95(self, pred, target):
        """
        计算95% Hausdorff距离 (HD95)
        
        Args:
            pred: 预测结果
            target: 真实标注
            
        Returns:
            hd95: 95% Hausdorff距离
        """
        pred = pred.astype(bool)
        target = target.astype(bool)
        
        # 获取边界点
        pred_points = np.argwhere(pred)
        target_points = np.argwhere(target)
        
        if len(pred_points) == 0 or len(target_points) == 0:
            return np.inf
        
        # 计算所有点对之间的距离
        distances_pred_to_target = []
        for p in pred_points:
            distances = np.sqrt(np.sum((target_points - p) ** 2, axis=1))
            distances_pred_to_target.append(np.min(distances))
        
        distances_target_to_pred = []
        for t in target_points:
            distances = np.sqrt(np.sum((pred_points - t) ** 2, axis=1))
            distances_target_to_pred.append(np.min(distances))
        
        # 合并所有距离
        all_distances = distances_pred_to_target + distances_target_to_pred
        
        # 计算95%分位数
        hd95 = np.percentile(all_distances, 95)
        
        return hd95
    
    def compute_all_metrics(self, pred, target):
        """
        计算所有评估指标
        
        Args:
            pred: 预测结果
            target: 真实标注
            
        Returns:
            metrics: 包含所有指标的字典
        """
        metrics = {
            'dice': self.dice_coefficient(pred, target),
            'jaccard': self.jaccard_index(pred, target),
            'ppv': self.positive_predictive_value(pred, target),
            'sensitivity': self.sensitivity(pred, target),
            'hd': self.hausdorff_distance(pred, target),
            'hd95': self.hausdorff_distance_95(pred, target)
        }
        
        return metrics
    
    def compute_metrics_batch(self, pred_batch, target_batch):
        """
        批量计算评估指标
        
        Args:
            pred_batch: 批量预测结果 (B, C, H, W)
            target_batch: 批量真实标注 (B, C, H, W)
            
        Returns:
            metrics_list: 每个样本的指标列表
            mean_metrics: 平均指标
        """
        batch_size = pred_batch.shape[0]
        metrics_list = []
        
        for i in range(batch_size):
            pred = pred_batch[i, 0].cpu().numpy()
            target = target_batch[i, 0].cpu().numpy()
            
            # 二值化预测结果
            pred = (pred > 0.5).astype(np.uint8)
            target = target.astype(np.uint8)
            
            metrics = self.compute_all_metrics(pred, target)
            metrics_list.append(metrics)
        
        # 计算平均指标
        mean_metrics = {}
        for key in metrics_list[0].keys():
            values = [m[key] for m in metrics_list if not np.isinf(m[key])]
            mean_metrics[key] = np.mean(values) if values else 0.0
        
        return metrics_list, mean_metrics


class DiceLoss(torch.nn.Module):
    """Dice损失函数"""
    
    def __init__(self, smooth=1.0):
        super(DiceLoss, self).__init__()
        self.smooth = smooth
        
    def forward(self, pred, target):
        """
        计算Dice损失
        
        Args:
            pred: 预测结果 (B, C, H, W)
            target: 真实标注 (B, C, H, W)
            
        Returns:
            loss: Dice损失
        """
        # 展平张量
        pred = pred.view(-1)
        target = target.view(-1)
        
        # 计算交集
        intersection = (pred * target).sum()
        
        # 计算Dice系数
        dice = (2.0 * intersection + self.smooth) / (pred.sum() + target.sum() + self.smooth)
        
        # 返回损失 (1 - Dice)
        return 1.0 - dice


class CombinedLoss(torch.nn.Module):
    """组合损失函数 (BCE + Dice)"""
    
    def __init__(self, bce_weight=0.5, dice_weight=0.5):
        super(CombinedLoss, self).__init__()
        self.bce_weight = bce_weight
        self.dice_weight = dice_weight
        self.bce_loss = torch.nn.BCEWithLogitsLoss()
        self.dice_loss = DiceLoss()
        
    def forward(self, pred, target):
        """
        计算组合损失
        
        Args:
            pred: 预测结果 (未经过sigmoid)
            target: 真实标注
            
        Returns:
            loss: 组合损失
        """
        # BCE损失
        bce = self.bce_loss(pred, target)
        
        # Dice损失 (需要先sigmoid)
        pred_sigmoid = torch.sigmoid(pred)
        dice = self.dice_loss(pred_sigmoid, target)
        
        # 组合损失
        loss = self.bce_weight * bce + self.dice_weight * dice
        
        return loss